{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "# Import packages\n",
    "# import sys\n",
    "# sys.path.append('/home/gpu/peng')\n",
    "# from hj_dataset_devkit import SupportedDataset, load_dataset, CoordinateSystem\n",
    "from __init__ import SupportedDataset, load_dataset, CoordinateSystem"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load dataset\n",
    "# nuscenes\n",
    "# dataset_path = '/media/gpu/sdb/datasets/nuscenes/v1.0-trainval/'\n",
    "# dataset = load_dataset(SupportedDataset.nuscenes, dataset_path)\n",
    "# hongjing\n",
    "# dataset_path = '/media/gpu/sdd/hj_joint_dataset'\n",
    "# dataset = load_dataset(SupportedDataset.hongjing, dataset_path)\n",
    "# argoverse2\n",
    "dataset_path = '/media/gpu/sdb/datasets/argoverse2/sensor/'\n",
    "dataset = load_dataset(SupportedDataset.argoverse2, dataset_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Cache / summary dataset\n",
    "# dataset.scenes[0].cache('../cache_test', force_clear=False)\n",
    "# dataset.scenes = dataset.scenes[:30]\n",
    "# stat = dataset.get_statistic()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Usage of dataset\n",
    "wrong_path = '/home/gpu/peng/'\n",
    "print(dataset.check_file_structure(wrong_path))\n",
    "# dataset.load_dataset(dataset_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Usage of scene\n",
    "scene = dataset.scenes[0]\n",
    "print(f'scene name: {scene.meta.name}')\n",
    "print(f'scene description: {scene.meta.description}')\n",
    "print(f'scene tags: {[i.value for i in scene.meta.tags]}')\n",
    "print(f'scene lidar_channels: {scene.meta.lidar_channels}')\n",
    "print(f'scene radar_channels: {scene.meta.radar_channels}')\n",
    "print(f'scene calibration keys: {scene.calib.keys()}')\n",
    "print(f'{len(scene.frames)} frames in scene {scene.meta.name}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Usage of frame\n",
    "frame = scene.frames[0]\n",
    "print(f'lidar sensors: {frame.lidars}')\n",
    "print(f'radar sensors: {frame.radars}')\n",
    "print(f'camera sensors: {frame.cameras}')\n",
    "print(f'location in world of frame: {frame.ego_pose}')\n",
    "print(f'all sensors: {frame.all_sensors()}')\n",
    "print(f'frame timestamp(ms): {frame.get_timestamp()}')\n",
    "print(f'{frame.cameras[-1]} timestamp(ms): {frame.get_timestamp(frame.cameras[-1])}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get lidar/radar cloud of frame\n",
    "coor = CoordinateSystem.world\n",
    "lidar = frame.lidars[0]\n",
    "lidar_cloud = frame.get_lidar_cloud(lidar, coordinate_system=coor)\n",
    "print(f'{lidar} cloud point number: {lidar_cloud.shape[0]}')\n",
    "\n",
    "if len(frame.radars) > 0:\n",
    "    radar = frame.radars[-1]\n",
    "    radar_cloud = frame.get_radar_cloud(radar, coor)\n",
    "    print(f'{radar} cloud point number: {radar_cloud.shape[0]}')\n",
    "    frame.matplot_on_bev([lidar, radar], plot_obj=False, plot_map=True, coordinate_system=coor)\n",
    "else:\n",
    "    frame.matplot_on_bev([lidar], plot_obj=False, plot_map=True, coordinate_system=coor)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get camera image of frame\n",
    "camera = frame.cameras[0]\n",
    "img = frame.get_camera_image(camera)\n",
    "print(f'image of {camera}')\n",
    "\n",
    "frame.matplot_on_camera(camera, plot_lidar=True, plot_radar=True, plot_obj=True, plot_map=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get obstacles\n",
    "coor = CoordinateSystem.world\n",
    "obstacles = frame.get_obstacles(coor)\n",
    "if obstacles is not None:\n",
    "    frame.matplot_on_bev(frame.lidars + frame.radars, plot_obj=True, plot_map=False, coordinate_system=coor)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get map info\n",
    "MAP_INFO_RANGE = -1\n",
    "map_info = frame.get_map_info(coor, roi=MAP_INFO_RANGE)\n",
    "if map_info is not None:\n",
    "    print(f'{len(map_info.lane_dividers)} lane dividers within {MAP_INFO_RANGE}m')\n",
    "    print(f'{len(map_info.road_edges)} road edges within {MAP_INFO_RANGE}m')\n",
    "    print(f'{len(map_info.road_markings)} road markings within {MAP_INFO_RANGE}m')\n",
    "    print(f'{len(map_info.traffic_signs)} traffic signs within {MAP_INFO_RANGE}m')\n",
    "    print(f'{len(map_info.traffic_lights)} traffic lights within {MAP_INFO_RANGE}m')\n",
    "    frame.matplot_on_bev(frame.lidars + frame.radars, plot_obj=False, plot_map=True, coordinate_system=coor)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get tracked objects of scene/frame\n",
    "ca_interval_s = 4\n",
    "print(f'assuming that the movement in {ca_interval_s}s can be modeled by constant-acceleration')\n",
    "frame_id = int(len(scene.frames) / 2)\n",
    "frame_objs = scene.get_obstacles_with_motion(ca_interval_s, frame_id)\n",
    "print(f'frame {frame_id} has {len(frame_objs)} tracked obstacles')\n",
    "\n",
    "frames_objs = scene.get_obstacles_with_motion(ca_interval_s)\n",
    "print(f'{len(frames_objs)} frames have tracked obstacles')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get tracks in scene\n",
    "# ca_interval_s = 0.4\n",
    "scene_tracks = scene.get_obstacle_tracks(ca_interval_s)\n",
    "print(f'{len(scene_tracks)} tracked obstacles in scene {scene.meta.name}')\n",
    "scene_tracks[-1].matplot()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "MayaTools",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.13"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "c53cdaddb2dc5d9f0a8ad856a75eb3f00885100b5d9f842a291a96ffa7cb8a1d"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
